#!/usr/bin/env python
import argparse
import torch
import mindspeed.megatron_adaptor
import sys

def load_model_state(model_path):
    """
    加载 .pt 文件，并返回内部的 state_dict。
    若文件中存在键 "model"，则返回 checkpoint['model']，否则直接返回整个checkpoint。
    """
    checkpoint = torch.load(model_path, map_location="cpu")
    if isinstance(checkpoint, dict) and "model" in checkpoint:
        return checkpoint["model"]
    return checkpoint

def compare_state_dicts(state1, state2, rtol=1e-5, atol=1e-8):
    """
    对比两个 state_dict 中的权重：
      - 检查两者的 key 是否一致；
      - 对于每个 key 检查 shape 是否匹配；
      - 使用 torch.allclose 检查数值是否近似相等。
    返回：True 表示全部匹配，否则返回 False。
    """
    keys1 = set(state1.keys())
    keys2 = set(state2.keys())
    all_ok = True

    if keys1 != keys2:
        missing_in_first = keys2 - keys1
        missing_in_second = keys1 - keys2
        if missing_in_first:
            print(f"在第一个模型中缺失的权重键: {missing_in_first}")
        if missing_in_second:
            print(f"在第二个模型中缺失的权重键: {missing_in_second}")
        all_ok = False

    common_keys = keys1 & keys2
    for key in common_keys:
        tensor1 = state1[key]
        tensor2 = state2[key]
        if tensor1.shape != tensor2.shape:
            print(f"键 {key} 的形状不匹配: {tensor1.shape} vs {tensor2.shape}")
            all_ok = False
            continue
        if not torch.allclose(tensor1, tensor2, rtol=rtol, atol=atol):
            diff = (tensor1 - tensor2).abs().max()
            print(f"键 {key} 的数值不匹配，最大差异: {diff}")
            all_ok = False
    return all_ok

def main():
    parser = argparse.ArgumentParser(
        description="对比两个 .pt 模型文件中的权重是否匹配（考虑浮点误差）"
    )
    parser.add_argument("--model1", type=str, required=True,
                        help="第一个 .pt 模型文件路径")
    parser.add_argument("--model2", type=str, required=True,
                        help="第二个 .pt 模型文件路径")
    parser.add_argument("--rtol", type=float, default=1e-5,
                        help="比较时的相对容差")
    parser.add_argument("--atol", type=float, default=1e-8,
                        help="比较时的绝对容差")
    args = parser.parse_args()

    print(f"加载模型1: {args.model1}")
    state1 = load_model_state(args.model1)
    print(f"加载模型2: {args.model2}")
    state2 = load_model_state(args.model2)

    print("开始比较两个模型的权重 ...")
    match = compare_state_dicts(state1, state2, rtol=args.rtol, atol=args.atol)
    if match:
        print("所有权重匹配！")
        sys.exit(0)
    else:
        print("存在不匹配的权重！")
        sys.exit(1)

if __name__ == "__main__":
    main()
